#!/usr/bin/env python

import numpy as np
from scipy.linalg import svd
import pickle, os, sys
sys.path.append('../data_parsers')
import ols
from encode_truth import encode_truth
from datetime import datetime
siteid = 'KBIL'


def main():
    compute_coefs()
    #evaluate_fcst()


def evaluate_fcst(mosdict,fcst_date):
    coeffs = pickle.load(open('%s_gfse_moscoef.pickle' % siteid.upper(),'r'))
    # Since the dictionaries are keyed the same way, we can just
    # grab the coefficient
    # Also need day of year
    doy = fcst_date.timetuple().tm_yday
    if doy > 183:
        doy = 366 - doy
    vdict = {}
    for var in coeffs.keys():
        # Do the constant and day of year first
        vdict[var] = coeffs[var]['const']
        vdict[var] = coeffs[var]['dayofyr'] * doy
        for coe in coeffs[var].keys():
            if coe not in ['const','dayofyr']:
                newval = coeffs[var][coe] * mosdict[coe]
                vdict[var] = vdict[var] + newval

    print "MOS FORECAST:"
    for var in vdict.keys():
        print var, vdict[var]
    return vdict




def compute_coefs():
    outdict = {}
    gfse_arch = pickle.load(open('%s_gfse_archive.pickle' % siteid.upper(),'r'))
    truth = encode_truth(siteid)

    # Now we can start working
    # Get a list of dates
    datelist = gfse_arch[0].keys()
    datelist.sort()

    # Four variables to predict
    predict_vars = ['high','low','wind','precip']
    # loop through each
    for pvar in predict_vars:
        print pvar
        truthlist = [truth[d][pvar] for d in datelist]
        print "NUMDATES:", len(truthlist)
        # Now make a list of each ensemble member's
        # output variables and make it a huge array
        thematrix = []
        varnames = []
        ensmems = gfse_arch.keys()
        ensmems.sort()
        for mem in ensmems:
            for v in ['high','low','dpt','cld','p24','q24']:
                vlist = [gfse_arch[mem][d][v] for d in datelist]
                thematrix.append(vlist)
                varnames.append('m%d_%s' % (mem,v))

        # Append the day of year
        daylist = [d.timetuple().tm_yday for d in datelist]
        daylist2 = [d if (d < 183) else (366-d) for d in daylist]
        thematrix.append(daylist2)
        varnames.append('dayofyr')

        thematrix = np.array(thematrix)
        thematrix = np.transpose(thematrix)
        print "MATRIX SHAPE", np.shape(thematrix)
        #print varnames
        #raw_input()
        # Use ols to run the model
        omodel = ols.ols(np.array(truthlist),thematrix,pvar,varnames)
        #print "pvalues:", omodel.p
        #print omodel.summary()
        #print dir(ols)
        #raw_input()
        coeffs = omodel.b
        # Save in a dictionary
        outdict[pvar] = {}
        outdict[pvar]['const'] = coeffs[0]
        for coef,vn in zip(coeffs[1:],varnames):
            outdict[pvar][vn] = coef
    pickle.dump(outdict,open('%s_gfse_moscoef.pickle' % siteid.upper(),'w'))
    

if __name__ == '__main__':
    main()
